%%% this code runs the counterfactuals

clear
close all
clc
tic
set(0,'DefaultFigureWindowStyle','docked','DefaultFigureVisible','on')

%% folder locations

folders;

%% set parameters
%calibrate constants?
calibrate_d0 = false;      % if true, only calibrated if the country is Spain
calibrate_d1 = false;

% parameters of grid to load
EGM = 8;

% parameters that remain constant throughout the loop across calibrations
alpha = 0.4;
beta = 0.13; %1.2446; 
sigma = 5;
rho = 0;
a = 1;
nu = 1;
NUTS = 2;  % must be = 2

% use adigator?
adigator = 0;

GAMMA = ( 0.10/0.13 )*beta;
MOBIL = {'off'}; % 'on'
CONGESTION = {'on'};

country='MERCOSUR'; % define group of connected countries to calibrate 'MERCOSUR' or 'ANDES'

%% loop
for cong = CONGESTION

    for gamma = GAMMA    

        for mobil = MOBIL

            % baseline parameters 
            % this is used to create filename
            % number of goods may be redefined below to determine Ngoods in 'nuts' case                

            param = init_parameters( 'a',a,'rho',rho,'alpha',alpha,'sigma',sigma,...       
                                     'beta',beta,'gamma',gamma,'nu',nu,...   
                                     'K',1,...
                                     'LaborMobility',char(mobil),...
                                     'CrossGoodCongestion',char(cong),...                                         
                                     'TolKappa',1e-4 );  

            % define filename to save diary, calibration and counterfactuals     
            filename = [ country,...
                        '_a',num2str( param.a ),...
                        '_rho',num2str( param.rho ),...
                        '_alpha',num2str( param.alpha ),...
                        '_sigma',num2str( param.sigma ),...
                        '_beta',num2str( param.beta ),...
                        '_gamma',num2str( param.gamma ),...
                        '_nu',num2str( param.nu ),...
                        '_mobil',num2str( param.mobility ),...
                        '_cong',num2str( param.cong )];

   
            % start diary
            diary( [path_save_calibrations,filename,'_calib.log'] ); % keep track of matlab output in log file
            clc   

            fprintf('----------------------------------------------------------------\n');
            fprintf('CALIBRATION - Country: %s, mobil = %d, cong = %d, gamma = %1.2f\n\n',country,param.mobility,param.cong,param.gamma);

            % calibrate d0 only if country is Spain and we run the calibration
            if strcmp(country,'Spain') && (calibrate_d0==true)                    
                calibrate_d0 = true;
                save_calibrated_d0 = 1;
            else
                calibrate_d0 = false;
                save_calibrated_d0 = 0;       
            end                

            % recover graph 
            graph_filename =  [ path_save_grids,country ];                       
            load( [ graph_filename, '_grid.mat' ] );

            % graph
            g = country_graph.graph_export;
            unique_edges = country_graph.unique_edges;
            gridmap = country_graph.gridmap;

            % Define goods and productivities 
            uniqueCountries = unique({gridmap(:).Country});
            Ncountries = length(uniqueCountries);
            Nrows_perCountry = zeros(Ncountries, 1);
            if strcmp(country, 'MERCOSUR')
                Ngoods_perCountry = {5, 7, 1, 1}; % number of DPs corresponding to country list abov
            
            elseif strcmp(country, 'ANDES')
                  Ngoods_perCountry =  {1, 3, 3, 3};
            end
            
            Ngoods = sum(cell2mat( Ngoods_perCountry))+1;

            % reinit param with the right # of goods -- need to calibrate this value
            param=init_parameters('param',param,'N',Ngoods,'m',ones(Ngoods,1)); % reinit param with the right # of goods            

            % graph_export contains information about the attributes of the discretized road network
            % graph_export.J = number of nodes
            % graph_export.nodes{n} = structure with fields specific to the node
            %    graph_export.nodes{n}.neighbors   = 1xN row with list of neighbors
            %    graph_export.nodes{n}.neighbors.x = x-coordinate
            %    graph_export.nodes{n}.neighbors.y = y-coordinate
            %    graph_export.adjacency = adjacency matrix (zero-ones)
            %    graph_export.avI = matrix with avI, equal to 1e-4 if nolink
            %    graph_export.distance = matrix with distance between connected nodes on the underlying graph
                    % note: these are distances on shortest path on the
                    % network if there is a link. if no link, distance is the geographic distance
            %   graph_export.L and graph_export.Y are Jx1 columns with population and income in each cell
            %   graph_export.av_alt = altitude
            %   graph_export.sd_alt = sd of altitude
            %   graph_export.rugged = ruggedness
            %   graph_export.all_distances = J*J matrix with bilateral geographic distances

            % we use the notation in the paper:
            %
            %       tau( Q,I ) = delta_tau * Q^beta * I^( -gamma )
            %

            % Fix some issues with data
            I = find( g.distance == 0 & g.adjacency==1 );
            g.distance(I) = mean(g.distance(:));

            %  parametrize delta_tau                
            %d0 = 1; % default value of d0
            d1 = 1;
            d0 = load_d0_from_file('calibration.dat',beta,gamma,param.mobility,param.cong,10 );
            g.delta_tau = d0*( g.distance ).^d1;          %  From P+E's calibrate_europe.m.            

            % ------------------
            % SET UP THE ECONOMY
            % ------------------

            % Complete and rearrange data structure
            g.ndeg = sum(reshape(tril(g.adjacency),[g.J^2,1])); % degrees of freedom
            g.Y = double(g.Y); % IPOPT only works with double precision
            g.L = double(g.L);
            g.Y = g.Y/sum(g.Y);   % normalize GDP to 1
            g.L = g.L/sum(g.L); % normalize population to 1        

            % Correct issues with data
            I=find(g.L==0); % identify places with no population
            g.Y(I) = 1e-5;  % set their production to 1e-5
            g.L(I) = 1e-5;  % and their population to 1e-5

            % retrieve number of locations
            param.J = g.J;

            % Define population and housing per capita
            if param.mobility==false
                param.Lj = g.L;
                param.hj=ones(g.J,1); % if not calibrated, housing is calibrated to 1 per capita
                param.Hj=param.hj.*param.Lj;
                param.omegaj=ones(g.J,1);
            else                        
                param.Hj=ones(g.J,1);
            end

            for m=1:length( uniqueCountries )
                Nrows_perCountry(m,1) = sum(strcmp({gridmap(:).Country},uniqueCountries(m)));
            end

            param.Zjn = zeros( g.J,Ngoods );
            param.Zjn(:,1) = 1; % all cities produce the agricultural good

            c=1;
            k=1;
            for z=1:length( uniqueCountries ) % loop through each country one at a time 
                n=1;
                locations= find(strcmp({gridmap(:).Country},uniqueCountries(z))); % locations for this country

                while n<=cell2mat(Ngoods_perCountry(z)) && ~isempty(locations)
                    
                    [Y,I]=sort(g.L(locations),'descend'); % sort by population
                    param.Zjn(locations(I(1)),1+k)=1; % only largest cities produce one of the specific goods
                    param.Zjn(locations(I(1)),1)=0;

                    % remove largest city and its neighbors from the list
                    locations=setdiff(locations,[locations(I(1)),g.nodes{locations(I(1))}.neighbors]);
                    n=n+1;
                    k=k+1;
                end
                c = Nrows_perCountry(m)+1;
            end
            
            param.N=Ngoods;
            param.Zjn = param.Zjn(:,1:param.N);
            param.m = param.m(1:param.N);

            %identify differentiated nodes
            nodes_dif = find(sum(param.Zjn(:,2:end)>0,2));

            %% calibrate
            param.adigator = adigator;                    
            calibrate(); 
            results_actual = results; % keep the calibrated allocation for welfare comparison
            results_actual.Ijk = g.avI;

            % DO NOT Compute additional statistics
            % results_actual.taxes = compute_taxes(param,g,results_actual);
            % results_actual.flaw_coeffs = compute_fundamental_law(param,g,results_actual);

            % Identify the delta_i_tilde's
            if param.cong==false % no cross-good congestion
                Pjkn=repmat(permute(results.Pjn,[1 3 2]),[1 g.J 1]);
                PQ=Pjkn.*results.Qjkn.^(1+param.beta);
                delta_i_tilde = g.delta_tau.*sum(PQ+permute(PQ,[2 1 3]),3)./(g.avI.^(1+param.gamma));
            else % cross-good congestion
                PCj=repmat(results.PCj,[1 g.J]);
                matm=shiftdim(repmat(param.m,[1 g.J g.J]),1);
                cost=sum(matm.*results.Qjkn.^param.nu,3).^((param.beta+1)/param.nu);
                PQ=PCj.*cost;
                delta_i_tilde = g.delta_tau.*(PQ+PQ')./(g.avI.^(1+param.gamma));
            end

            delta_i_tilde(~g.adjacency)=0;     
            g.delta_i_tilde = delta_i_tilde;
            g.delta_i_tilde = g.delta_i_tilde / sum( g.delta_i_tilde(:).*g.avI(:) );  % normalize to meet constraint with K=1

            %% project delta_i_tilde on fundamentals       

            % construct dataset at bilateral level
            id=find( tril(g.adjacency)==1 );
            Npairs = length(unique_edges);  % number of unique (unordered) origin-destination pairs

            d_i_tilde = zeros( Npairs,1 );
            dist = zeros( Npairs,1 );
            diff_alt = zeros( Npairs,1 );
            av_alt = zeros( Npairs,1 );
            av_rugged = zeros( Npairs,1 );
            %FE = zeros( Npairs,g.J );  % fixed effect

            data = cell( Npairs,10 );

            for i=1:Npairs   % each row is a pair

                node1 = unique_edges( i,1 );
                node2 = unique_edges( i,2 );

                d_i_tilde( i ) = log( delta_i_tilde( node1,node2 ) );
                dist( i ) = log( g.distance( node1,node2 ) );
                diff_alt( i ) = abs( log( g.av_alt( node1 )/g.av_alt( node2 ) ) );
                av_alt( i ) = 1/2*( log( g.av_alt( node1 ) )+log( g.av_alt( node2 ) ) );
                av_rugged( i ) = 1/2*( log( g.rugged( node1 ) )+log( g.rugged( node2 ) ) );

                % matrix to export to csv
                data{i,1} = country;
                data{i,2} = country; 
                data{i,3} = node1;
                data{i,4} = node2;
                data{i,5} = delta_i_tilde( node1,node2 );
                data{i,6} = g.distance( node1,node2 );
                data{i,7} = g.av_alt( node1 );
                data{i,8} = g.av_alt( node2 );
                data{i,9} = g.rugged( node1 );
                data{i,10} = g.rugged( node2 );

            end

            % run linear model       
            tbl = table( d_i_tilde,dist,diff_alt,av_alt,av_rugged,...
                         'VariableNames',{'DeltaTilde','Distance','Diff_Alt','Av_Alt','Av_Rugged'} );                 
            mdl = fitlm( tbl,'DeltaTilde~Distance+Diff_Alt+Av_Rugged' );             

            % project delta_i
            g.delta_i_projected = zeros(g.J,g.J);
            for k=1:Npairs
                g.delta_i_projected(unique_edges(k,1),unique_edges(k,2)) = exp(mdl.Fitted(k));
            end

            g.delta_i_projected=g.delta_i_projected+g.delta_i_projected';  
            g.delta_i_projected( g.adjacency == 0 ) = 0;
            g.delta_i_projected = g.delta_i_projected / sum( g.delta_i_projected(:).*g.avI(:) );  % normalize to meet constraint with K=1                

            %% Compute also the delta_i_engineer using Collier's formula
            g.edge_ruggedness = repmat(g.rugged,[1 g.J]).^.5.*repmat(g.rugged',[g.J 1]).^.5;
            g.edge_ruggedness( g.adjacency == 0 )=0;
      
            g.delta_i_engineer = exp( -0.11*( g.distance > 50) + ...
            0.12*log( g.edge_ruggedness ) + log( g.distance ) );

            g.delta_i_engineer( g.adjacency == 0 ) = 0;
            g.delta_i_engineer = g.delta_i_engineer / sum( g.delta_i_engineer(:).*g.avI(:) );  % normalize to meet constraint with K=1

            %% save calibrated model

            calibration.g = g; % note: the deltas saved under g are not normalized
            calibration.param = param;
            calibration.results_actual = results_actual;
            
            save( [ path_save_calibrations,filename,'_collier_calib.mat' ],'calibration' );                       
   
            % terminate diary (saved under calibrations folder)
            diary off;

        end
    end
end
toc


